import equinox as eqx
from jaxtyping import Array, Float, Int, Bool, PyTree
import jax.numpy as jnp
import jax
from .modules.utils import forward_pass, get_dts
from .modules.ode import integrate
import math

class Func(eqx.Module):
    net: list
    intensity_fn: list

    def __init__(self, hdim: int, num_types: int, key: Array):
        key1, key2, key3, key4 = jax.random.split(key, 4)
        self.net = [
            eqx.nn.Linear(hdim, 2*hdim, key=key1),
            jax.nn.silu,
            eqx.nn.Linear(2*hdim, hdim, key=key2),
            jax.nn.tanh
        ]
        self.intensity_fn = [
            eqx.nn.Linear(hdim, 2*hdim, key=key3),
            jax.nn.silu,
            eqx.nn.Linear(2*hdim, num_types, key=key4),
            jax.nn.softplus
        ]

    def __call__(self, t: Float[Array, ""], state: PyTree, args=None):
        energy, Lambda, z = state
        dz = forward_pass(self.net, z)
        intensities = self.get_intensities(z)
        intensity = intensities.sum(-1)
        denergy = ((dz**2).sum()) / (dz.size)
        return denergy, intensity, dz
        
    def get_intensities(self, z: Float[Array, "hdim"]):
        intensities = forward_pass(self.intensity_fn, z)
        return intensities
    
class Jump(eqx.Module):
    net: list

    def __init__(self, hdim: int, embed_dim: int, key: Array):
        key1, key2 = jax.random.split(key)
        self.net = [
            eqx.nn.Linear(hdim+embed_dim, 2*hdim, key=key1),
            jax.nn.silu,
            eqx.nn.Linear(2*hdim, hdim, key=key2)
        ]

    def __call__(self, embed: Float[Array, "embed_dim"], z: Float[Array, "hdim"]):
        z_ = forward_pass(self.net, jnp.concatenate((z, embed), -1))
        return z_
    
class Encoder(eqx.Module):
    func: Func
    jump: Jump
    init_state: Array

    def __init__(self, hdim: int, embed_dim: int, num_types: int, key: Array):
        key1, key2, key3 = jax.random.split(key, 3)
        self.func = Func(hdim, num_types, key1)
        self.jump = Jump(hdim, embed_dim, key2)
        lim = 1 / math.sqrt(hdim)
        self.init_state = jax.random.uniform(key3, (hdim,), minval=-lim, maxval=lim)

    def extrapolate(self, state0: PyTree, dt: Float[Array, ""]):
        state1 = integrate(self.func, 0., dt, state0, None)
        return state1
    
    def step(self, state0: PyTree, dt: Float[Array, ""], embed: Float[Array, "hdim"]):
        state1 = integrate(self.func, 0., dt, state0, None)
        energy1, Lambda1, z1_prior = state1
        z1_post = self.jump(embed, z1_prior)
        return energy1, Lambda1, z1_post, z1_prior
    
    def __call__(self, dts: Float[Array, "T"], embeds: Int[Array, "T hdim"] ):
        state0 = (jnp.asarray(0.), jnp.asarray(0.), self.init_state)  # (energy, Lambda, Z)
        inputs = (dts, embeds)
        def step(h, x):
            dt, embed = x
            energy1, Lambda1, z1_post, z1_prior = self.step(h, dt, embed)
            carry = (energy1, Lambda1, z1_post)
            y = (z1_post, z1_prior)
            return carry, y
        carry, ys = jax.lax.scan(step, state0, inputs)
        energy, Lambda, z_end = carry
        zs_post, zs_prior = ys
        return zs_post, zs_prior, energy, Lambda, z_end

class ODETPP(eqx.Module):
    enc: Encoder
    energy_reg: float
    embed: eqx.nn.Embedding

    def __init__(self, hdim: int, embed_dim: int, num_types: int, energy_reg: float, key: Array):
        key1, key2 = jax.random.split(key)
        self.enc = Encoder(hdim, embed_dim, num_types, key1)
        self.embed = eqx.nn.Embedding(num_types, embed_dim, key=key2)
        self.energy_reg = energy_reg

    def __process_Z(self, z: Float[Array, "hdim"], mark: Int[Array, ""]):
        intensities = self.enc.func.get_intensities(z)
        intensity = intensities[mark]
        log_intensity = jnp.log(intensity + 1e-8)
        mark_prob = intensity / (intensities.sum()+1e-8)
        mark_ll = jnp.log(mark_prob+1e-8)
        return log_intensity, mark_ll

    @eqx.filter_jit
    def __call__(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.) # (T)
        embeds = jax.vmap(self.embed)(marks)
        zs_post, zs_prior, energy, Lambda, z_end = self.enc(dts, embeds)
        log_intensity, mark_ll = jax.vmap(self.__process_Z)(zs_prior[1:], marks[1:])
        mark_ll = jnp.where(mask[1:], mark_ll, 0.)  # (T-1)
        log_intensity = jnp.where(mask[1:], log_intensity, 0.)
        mark_ll = mark_ll.sum()
        ll = log_intensity.sum() - Lambda
        time_ll = ll - mark_ll
        return ll, time_ll, mark_ll, energy* self.energy_reg

    def _predict(self, z: Float[Array, "hdim"], dt_max: Float[Array, ""]):
        energy0 = jnp.asarray(0.)
        Lambda0 = jnp.asarray(0.)
        state0 = (energy0, Lambda0, z)
        int_tf0 = jnp.asarray(0.)
        H0 = (state0, int_tf0)
        def func(t, H, args=None):
            state, int_tf = H
            energy, Lambda, z = state
            d_state = self.enc.func(t, state)
            d_energy, intensity, dz = d_state
            f_dt = intensity * jnp.exp(-Lambda)
            tf = t * f_dt
            return d_state, tf
        _, Efdt = integrate(func, 0., dt_max, H0, None)
        Efdt = Efdt.clip(min=0.)
        state1 = self.enc.extrapolate(state0, Efdt)
        energy1, Lambda1, z1 = state1
        intensities = self.enc.func.get_intensities(z1) # (num_types, )
        mark_predict = intensities.argmax(-1)
        return Efdt, mark_predict
    
    @eqx.filter_jit
    def rolling_predict(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], dt_max: float, key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.) # (T)
        embeds = jax.vmap(self.embed)(marks)
        zs_post, zs_prior, energy, Lambda, z_end = self.enc(dts, embeds)
        dt_predict, mark_predict = jax.vmap(self._predict, (0, None))(zs_post[:-1], dt_max)  # (T-1, ), (T-1, )
        return (dt_predict, mark_predict), (dts[1:], marks[1:]), mask[1:]